from openmind_hub import create_repo,create_branch,upload_file,upload_folder,snapshot_download,om_hub_download,OM_HUB_CACHE
import unittest
from datetime import date
import time
import os,string,random
import shutil

token = "123456789()"
repo_id = "xiaoguozhi/gxz"+date.today().strftime('%Y%m%d')
branch = "v1"
author = "xiaoguozhi"
reponame = "gxz"+date.today().strftime('%Y%m%d')
def search_dir_File(pathname,reponame):
    folders = [f for f in os.listdir(pathname) if os.path.isdir(os.path.join(pathname, f))]
    filtered_folders = [f for f in folders if 'openmind' in f or 'model' in f or 'space' in f or 'dataset' in f or reponame in f]
    return filtered_folders

class TestOpenmindHub(unittest.TestCase):
    def test_01_hf_create_modelrepo_success(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                token=token,
                exist_ok=False,
            )
        except:
            raise Exception("创建仓库有报错")
    def test_02_hf_create_modelrepo_duplicate_existok(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                token=token,
                repo_type="model",
                exist_ok=True,
            )
        except:
            raise Exception("exist_ok=True未生效")
    def test_03_hf_create_modelrepo_duplicate_existnotok(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                token=token,
                repo_type="model",
                exist_ok=False,
            )
            raise Exception("exist_ok=False未生效")
        except:
            pass
    def test_04_hf_create_modelrepo_no_repo_type(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                token=token,
                exist_ok=True,
            )
        except:
            raise Exception("仓库类型未添加默认值")
    def test_05_hf_create_modelrepo_wrong_token(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                token=token+"a",
                repo_type="model",
                exist_ok=False,
            )
            raise Exception("token未校验")
        except:
            pass
    def test_06_hf_create_modelbranch_success(self):
        try:
            create_branch(
                repo_id=repo_id,
                branch=branch,
                revision="main",
                token=token,
                repo_type="model",
                exist_ok=False
            )
        except:
            raise Exception("创建仓库分支有报错")
    def test_07_hf_create_modelbranch_duplicate_existok(self):
        try:
            create_branch(
                repo_id=repo_id,
                branch=branch,
                revision="main",
                token=token,
                repo_type="model",
                exist_ok=True
            )
        except:
            raise Exception("exist_ok=True未生效")
    def test_08_hf_create_modelbranch_duplicate_existnotok(self):
        try:
            create_branch(
                repo_id=repo_id,
                branch=branch,
                revision="main",
                token=token,
                repo_type="model",
                exist_ok=False
            )
            raise Exception("exist_ok=False未生效")
        except:
            pass
    def test_09_hf_create_modelbranch_noexistrepoid(self):
        try:
            create_branch(
                repo_id=repo_id+"a",
                branch="test",
                revision="main",
                token=token,
                repo_type="model",
                exist_ok=False
            )
            raise Exception("不存在的仓库创建分支未报错")
        except :
            pass
    def test_10_hf_create_modelbranch_wrong_token(self):
        try:
            create_branch(
                repo_id=repo_id,
                branch="test",
                revision="main",
                token=token+"a",
                repo_type="model",
                exist_ok=False
            )
            raise Exception("token未校验")
        except:
            pass
    def test_11_hf_create_modelbranch_no_repo_type(self):
        try:
            create_branch(
                repo_id=repo_id,
                branch=branch,
                revision="main",
                token=token,
                exist_ok=True
            )
        except:
            raise Exception("仓库类型未添加默认值")
    def test_12_hf_uploadfile_omhubdownload_success(self):
        try:
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/app.py",
                path_in_repo="ddd.py",
                repo_id=repo_id,
                token=token,
                revision="v1",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            om_hub_download(
                repo_id=repo_id,
                filename="ddd.py",
                subfolder=None,
                user_agent=None,
                local_dir="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                force_download=False,
                token=token,
                local_files_only=False,
                revision="v1"
            )
            if os.path.exists("/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/ddd.py"):
                print("文件存在。")
            else:
                raise AssertionError("文件不存在。")
        except:
            raise Exception("下载文件有报错")
        else:
            os.remove("/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/ddd.py")
    def test_13_hf_uploadfile_noexistropo(self):
        try:
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub\app.py",
                path_in_repo="ddd.py",
                repo_id=repo_id+"A",
                token=token,
                revision="v1",
                commit_message="upload file",
                commit_description=None,
            )
            raise Exception("上传文件到不存在的仓库未报错")
        except:
            pass
    def test_14_hf_uploadfile_wrongtoken(self):
        try:
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub\app.py",
                path_in_repo="ddd.py",
                repo_id=repo_id,
                token=token+"A",
                revision="v1",
                commit_message="upload file",
                commit_description=None,
            )
            raise Exception("上传文件未校验token")
        except:
            pass
    def test_15_hf_uploadfile_duplicate(self):
        try:
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/app.py",
                path_in_repo="ddd.py",
                repo_id=repo_id,
                token=token,
                revision="v1",
                commit_message="upload file",
                commit_description=None,
            )
        except :
            raise Exception("重复上传文件报错")
    def test_16_hf_upload_folder_success(self):
        file_dir = []
        # 创建4000个随机文件
        for i in range(10):
            # 生成随机文件名
            random_file_name = ''.join(random.choices(string.ascii_letters + string.digits, k=6))
            file_path = f"/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/{random_file_name}.txt"
            file_dir.append(file_path)
            # 创建空白文件
            with open(file_path, 'w'):
                pass
        print("已创建4000个随机文件")
        try:
            upload_folder(
                repo_id=repo_id,
                folder_path="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                path_in_repo="/filedata",
                commit_message="upload folder",
                commit_description=None,
                token=token,
                revision=None,
                # allow_patterns=["*.bin", "*.py"],
                ignore_patterns="*.log",
                # num_threads=5,
            )
        except:
            raise Exception("上传文件报错")
        finally:
            for i in file_dir:
                os.remove(i)
    def test_17_hf_om_hub_download_wrongtoken(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="ddd.py",
                subfolder=None,
                user_agent=None,
                local_dir="/openmind-hub/auto_test/autotest/test_openmind_hubfiledata",
                force_download=False,
                token=token + "a",
                local_files_only=False,
                revision="v1"
            )
            raise Exception("omhubdownload未校验token")
        except:
           pass
    def test_18_hf_om_hub_download_noexistrepo(self):
        try:
            om_hub_download(
                repo_id=repo_id + "a",
                filename="ddd.py",
                subfolder=None,
                user_agent=None,
                local_dir="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                force_download=False,
                token=token,
                local_files_only=False,
                revision="v1"
            )
            raise Exception("omhubdownload未校验repo")
        except:
           pass
    def test_19_hf_snapshot_download_cachedir_success(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                cache_dir=OM_HUB_CACHE,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(OM_HUB_CACHE,"")
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
    def test_20_hf_snapshot_download_wrong_token(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                cache_dir=OM_HUB_CACHE,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token+"a",
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("下载文件未校验token")
        except :
           pass
    def test_21_hf_snapshot_download_noexist_repo(self):
        try:
            snapshot_download(
                repo_id=repo_id+"a",
                repo_type=None,
                revision=None,
                cache_dir=OM_HUB_CACHE,
                #local_dir="space_test0516",
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("下载文件未校验repo")
        except :
           pass
    def test_22_hf_snapshot_download_localdir_success(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type="model",
                revision="main",
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"/openmind-hub/auto_test/autotest/test_openmind_hub/" + repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs = search_dir_File(r"/openmind-hub/auto_test/autotest/test_openmind_hub/" + author, reponame)
            self.assertTrue(len(dirs), "下载的文件不存在")
        except:
            raise Exception("下载文件有报错")
        else:
            shutil.rmtree(r"/openmind-hub/auto_test/autotest/test_openmind_hub/" + author)
    def test_23_hf_snapshot_download_type_id_notconsistent(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type="dataset",
                revision="main",
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"/openmind-hub/auto_test/autotest/test_openmind_hub/" + repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("没有校验type与ID不一致的情况")
        except:
           pass





if __name__ == '__main__':
    unittest.main()